import os
import pathlib
import numpy as np
import tensorflow as tf
import scipy.io.wavfile as wav
from collections import *
import glob
import random

seed = 42
tf.random.set_seed(seed)
np.random.seed(seed)
# DATASET_PATH = '/mnt/VOICE/Data/mini_speech_commands'
DATASET_PATH="/mnt/VOICE/Data/kewos/kws_trans/training"
data_dir = pathlib.Path(DATASET_PATH)
commands = np.array(os.listdir(data_dir))
commands = commands[commands != 'README.md']
nums_class = len(commands)
commands2id ={}
for k,v in enumerate(commands):
    commands2id[v]=k

def decode_audio(file_path, seq_len):
    sample_rate, ori_siginal = wav.read(file_path)
    ori_siginal = list(ori_siginal)
    if len(ori_siginal) > seq_len:
        ori_siginal = ori_siginal[:seq_len]
    else:
        for i in range(seq_len - len(ori_siginal)):
            ori_siginal.append(0)
    return ori_siginal

# train_filenames = glob.glob("str(data_dir) + '/*/*'")
train_filenames = glob.glob("/mnt/VOICE/Data/kewos/kws_trans_leaf/training/*/*")
test_filenames =glob.glob("/mnt/VOICE/Data/kewos/kws_trans_leaf/testing/*/*")
val_filenames =glob.glob("/mnt/VOICE/Data/kewos/kws_trans_leaf/validation/*/*")
random.shuffle(train_filenames)
random.shuffle(test_filenames)
random.shuffle(val_filenames)

train_files = train_filenames
val_files = val_filenames
test_files = test_filenames

def get_label(file_path):
  parts = file_path.split(sep=os.path.sep)
  return commands2id[parts[-2]]
def decode_audio(audio_binary):
    # Decode WAV-encoded audio files to `float32` tensors, normalized
    # to the [-1.0, 1.0] range. Return `float32` audio and a sample rate.
    audio, _ = tf.audio.decode_wav(contents=audio_binary)
    audio = tf.cast(audio,tf.float32)
    # Since all the data is single channel (mono), drop the `channels`
    # axis from the array.
    return tf.squeeze(audio, axis=-1)


def get_label(file_path):
    parts = tf.strings.split(
        input=file_path,
        sep=os.path.sep)
    # Note: You'll use indexing here instead of tuple unpacking to enable this
    # to work in a TensorFlow graph.
    label = 0
    for i, j in enumerate(commands):
        if j == parts[-2]:
            label = i
    return label


def get_waveform_and_label(file_path):
    label = get_label(file_path)
    audio_binary = tf.io.read_file(file_path)
    print(audio_binary)
    waveform = decode_audio(audio_binary)

    return {"audio": tf.convert_to_tensor(waveform), "label": tf.convert_to_tensor(label),'file_name':tf.convert_to_tensor(file_path)}



AUTOTUNE = tf.data.experimental.AUTOTUNE
train_files_ds = tf.data.Dataset.from_tensor_slices(train_files)
val_files_ds = tf.data.Dataset.from_tensor_slices(val_files)
test_files_ds = tf.data.Dataset.from_tensor_slices(test_files)
train_waveform_ds = train_files_ds.map(
    map_func=get_waveform_and_label,
    num_parallel_calls=AUTOTUNE)

val_waveform_ds = val_files_ds.map(
    map_func=get_waveform_and_label,
    num_parallel_calls=AUTOTUNE)

test_waveform_ds = test_files_ds.map(
    map_func=get_waveform_and_label,
    num_parallel_calls=AUTOTUNE)

all_ds = {"train":train_waveform_ds,"validation":val_waveform_ds,"test":test_waveform_ds}

if __name__=="__main__":
    for i in all_ds["train"]:
        print(i['file_name'].numpy())

